{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e86b9824",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Users\\a\\.conda\\envs\\langchainV1\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from langchain.chat_models import init_chat_model\n",
    "from langchain_community.utilities import SQLDatabase\n",
    "from langchain.agents import create_agent\n",
    "from langchain_core.messages import SystemMessage\n",
    "from langgraph.checkpoint.memory import InMemorySaver\n",
    "import re, os\n",
    "from langchain_core.tools import tool\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "# 初始化模型\n",
    "llm = init_chat_model(\n",
    "    \"deepseek:deepseek-chat\",\n",
    "    temperature=0\n",
    "    )\n",
    "\n",
    "# mysql数据库配置\n",
    "db_host = os.getenv(\"DB_HOST\")\n",
    "db_port = os.getenv(\"DB_PORT\")\n",
    "db_user = os.getenv(\"DB_USER\")\n",
    "db_password = os.getenv(\"DB_PASSWORD\")\n",
    "db_name = os.getenv(\"DB_NAME\")\n",
    "\n",
    "MYSQL_URI = \"mysql+pymysql://{}:{}@{}:{}/{}\".format(\n",
    "    db_user, db_password, db_host, db_port, db_name\n",
    ")\n",
    "\n",
    "# 连接到MySQL数据库\n",
    "db = SQLDatabase.from_uri(MYSQL_URI)  \n",
    "# 获取数据库表的定义信息\n",
    "SCHEMA = db.get_table_info()\n",
    "\n",
    "# 用于检测和阻止危险的 SQL 操作\n",
    "DENY_RE = re.compile(r\"\\b(INSERT|UPDATE|DELETE|ALTER|DROP|CREATE|REPLACE|TRUNCATE)\\b\", re.I)\n",
    "# 检查 SQL 查询末尾是否已经包含 LIMIT 子句\n",
    "HAS_LIMIT_TAIL_RE = re.compile(r\"(?is)\\blimit\\b\\s+\\d+(\\s*,\\s*\\d+)?\\s*;?\\s*$\")\n",
    "\n",
    "def _safe_sql(q: str) -> str:\n",
    "    # 规范化查询\n",
    "    q = q.strip()\n",
    "    # 阻止多个语句（允许一个可选的尾随;）\n",
    "    if q.count(\";\") > 1 or (q.endswith(\";\") and \";\" in q[:-1]):\n",
    "        return \"错误：不允许多个语句。\"\n",
    "    q = q.rstrip(\";\").strip()\n",
    "\n",
    "    # 只读检查\n",
    "    if not q.lower().startswith(\"select\"):\n",
    "        return \"错误：只允许SELECT语句。\"\n",
    "    if DENY_RE.search(q):\n",
    "        return \"错误：检测到DML/DDL操作。只允许只读查询。\"\n",
    "\n",
    "    # 如果末尾没有LIMIT则添加（对空格/换行符鲁棒）\n",
    "    if not HAS_LIMIT_TAIL_RE.search(q):\n",
    "        q += \" LIMIT 5\"\n",
    "    return q\n",
    "\n",
    "@tool\n",
    "def execute_sql(query: str) -> str:\n",
    "    \"\"\"执行只读的MySQL SELECT查询并返回结果\"\"\"\n",
    "    query = _safe_sql(query)\n",
    "    q = query\n",
    "    if q.startswith(\"错误：\"):\n",
    "        return q\n",
    "    try:\n",
    "        return db.run(q)\n",
    "    except Exception as e:\n",
    "        return f\"错误：{e}\"\n",
    "    \n",
    "# 系统提示词\n",
    "SYSTEM = f\"\"\"您是一位谨慎的MySQL分析师。\n",
    "\n",
    "数据库表的定义信息（请勿虚构列/表）：\n",
    "{SCHEMA}\n",
    "\n",
    "规则：\n",
    "- 逐步思考。\n",
    "- 当需要数据时，使用一个SELECT查询调用工具`execute_sql`。\n",
    "- 仅限只读操作；不允许INSERT/UPDATE/DELETE/ALTER/DROP/CREATE/REPLACE/TRUNCATE。\n",
    "- 除非用户明确要求，否则限制为5行。\n",
    "- 如果工具返回'错误：'，请修改SQL并重试。\n",
    "- 尝试次数限制为5次。\n",
    "- 如果5次尝试后仍未成功，请向用户返回说明。\n",
    "- 优先使用明确的列列表；避免使用SELECT *。\n",
    "- 用户问题和数据库信息无关的情况下，直接回复用户，不要调用工具。\n",
    "\"\"\"\n",
    "\n",
    "# 创建内存存储器\n",
    "checkpointer = InMemorySaver()\n",
    "\n",
    "\n",
    "# 创建代理\n",
    "agent = create_agent(\n",
    "    model=llm,\n",
    "    tools=[execute_sql],\n",
    "    prompt=SystemMessage(content=SYSTEM),\n",
    "    checkpointer=checkpointer\n",
    ")\n",
    "\n",
    "config = {\"configurable\": {\"thread_id\": \"1\"}}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "225fa015",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "目前有多少名员工？\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "我来查询员工总数。\n",
      "Tool Calls:\n",
      "  execute_sql (call_00_FL60FzQyqW63r471YYB56Ndk)\n",
      " Call ID: call_00_FL60FzQyqW63r471YYB56Ndk\n",
      "  Args:\n",
      "    query: SELECT COUNT(*) as total_employees FROM employees\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: execute_sql\n",
      "\n",
      "[(7,)]\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "目前公司总共有7名员工。\n"
     ]
    }
   ],
   "source": [
    "question = \"目前有多少名员工？\"\n",
    "\n",
    "for step in agent.stream(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": question}]},\n",
    "    stream_mode=\"values\",\n",
    "    config=config,\n",
    "):\n",
    "    step[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cd2e0d5b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "各个部门的平均工资是多少？\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "我来查询各个部门的平均工资。首先需要关联员工表、部门关联表和部门表来计算每个部门的平均工资。\n",
      "Tool Calls:\n",
      "  execute_sql (call_00_kAWWhIuoqt4FwHQsylIP8Mwx)\n",
      " Call ID: call_00_kAWWhIuoqt4FwHQsylIP8Mwx\n",
      "  Args:\n",
      "    query: SELECT d.dept_name, AVG(e.salary) as avg_salary \n",
      "FROM departments d \n",
      "JOIN emp_dept_relation edr ON d.dept_id = edr.dept_id \n",
      "JOIN employees e ON edr.emp_id = e.emp_id \n",
      "GROUP BY d.dept_id, d.dept_name \n",
      "ORDER BY avg_salary DESC\n",
      "=================================\u001b[1m Tool Message \u001b[0m=================================\n",
      "Name: execute_sql\n",
      "\n",
      "[('技术部', Decimal('22500.000000')), ('市场部', Decimal('18500.000000')), ('财务部', Decimal('18000.000000')), ('人力资源部', Decimal('12000.000000'))]\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "各个部门的平均工资如下：\n",
      "\n",
      "1. **技术部**：22,500.00元\n",
      "2. **市场部**：18,500.00元  \n",
      "3. **财务部**：18,000.00元\n",
      "4. **人力资源部**：12,000.00元\n",
      "\n",
      "技术部的平均工资最高，人力资源部的平均工资相对较低。\n"
     ]
    }
   ],
   "source": [
    "question = \"各个部门的平均工资是多少？\"\n",
    "\n",
    "for step in agent.stream(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": question}]},\n",
    "    stream_mode=\"values\",\n",
    "    config=config,\n",
    "):\n",
    "    step[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fa6b1b98",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "================================\u001b[1m Human Message \u001b[0m=================================\n",
      "\n",
      "你好，我是tom\n",
      "==================================\u001b[1m Ai Message \u001b[0m==================================\n",
      "\n",
      "你好Tom！我是您的MySQL数据分析助手。我可以帮助您查询和分析数据库中的信息。\n",
      "\n",
      "我看到您的数据库包含以下表：\n",
      "- departments（部门信息表）\n",
      "- employees（员工信息表） \n",
      "- emp_dept_relation（员工部门关联表）\n",
      "\n",
      "如果您需要查询员工信息、部门信息或相关的数据分析，请告诉我您的具体需求，我会帮您查询相应的数据。\n"
     ]
    }
   ],
   "source": [
    "config = {\"configurable\": {\"thread_id\": \"2\"}}\n",
    "question = \"你好，我是tom\"\n",
    "\n",
    "for step in agent.stream(\n",
    "    {\"messages\": [{\"role\": \"user\", \"content\": question}]},\n",
    "    stream_mode=\"values\",\n",
    "    config=config\n",
    "):\n",
    "    step[\"messages\"][-1].pretty_print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e8c10bd",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchainV1",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
